from util.buffer import Buffer
from threading import RLock, Condition


class ObjectsFunc:
    def __init__(self, func, objs_limit=10):
        self.func = func
        self.buffer = Buffer(0)
        self.objs_limit = objs_limit

        self._result_ok = Condition(RLock())
        self._result_num = {}
        self._results = []
        self._id = 0

    def start(self):
        self.buffer.start()

    def stop(self):
        self.buffer.stop()

    def pause(self):
        self.buffer.pause()

    def resume(self):
        self.buffer.resume()

    def __call__(self, objs):
        with self._result_ok:
            id = self._id
            self._id += 1
            self._result_num[id] = len(objs)
        for i, obj in enumerate(objs):
            self.buffer.put((id, i, obj))
        return self._get_result(id)

    def _get_result(self, call_id):
        total = self._result_num[call_id]
        id_results = [None] * total
        completed = 0
        while completed < total:
            with self._result_ok:
                triples = self.buffer.get_many(empty_ok=True, limit=self.objs_limit)
                num = len(self._results)
                waiting = True
                if len(triples) > 0:
                    waiting = False
                    completed += self._extract_results(call_id, triples, id_results)
                for i in range(num-1, -1, -1):
                    id, index, result = self._result[i]
                    if id == call_id:
                        waiting = False
                        del self._results[i]
                        id_results[index] = result
                        completed += 1
                if waiting:
                    self._result_ok.wait()

        return id_results

    def _extract_results(self, call_id, triples, id_results):
        num = 0
        objs = []
        ids = []
        indices = []
        for id, index, obj in triples:
            ids.append(id)
            indices.append(index)
            objs.append(obj)
            results = self.func(objs)
        notify = False
        for id, index, result in zip(ids, indices, results):
            if id == call_id:
                id_results[index] = result
                num += 1
            else:
                self._results.append((id, index, result))
                notify = True
        if notify:
            self._result_ok.notify_all()
        return num


if __name__ == '__main__':
    def func(ints):
        return [i + 1 for i in ints]

    sf = ObjectsFunc(func, 3)

    from threading import Thread
    from numpy import random
    import time

    class Th(Thread):
        def __init__(self, id):
            super(Th, self).__init__(daemon=True)
            self.id = id

        def run(self):
            for _ in range(3):  # while True:
                ints = random.randint(1, 10)
                ints = [self.id * 10000 + i * 100 for i in range(ints)]
                print(ints, sf(ints))


    for i in range(20):
        th = Th(i)
        th.start()

    time.sleep(10)